# -*- coding: utf-8 -*-
# @Time    : 2024/10/11 11:39
# @Author  : yujiahao
# @File    : ag_dwd_insert.py
# @description: 异步执行SQL文件

import os
import sys
import asyncio
import asyncpg
import glob
import logging
import boto3
from botocore.exceptions import ClientError
import requests
import datetime
from typing import List, Dict, Set, Any
import sqlparse
from contextlib import asynccontextmanager

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 数据库连接配置，包含 dev、prod、local 三个环境
DB_CONFIGS = {
    'dev': {
        'database': os.getenv('DEV_REDSHIFT_DB'),
        'user': os.getenv('DEV_REDSHIFT_USER'),
        'password': os.getenv('DEV_REDSHIFT_PASSWORD'),
        'host': os.getenv('DEV_REDSHIFT_HOST'),
        'port': int(os.getenv('DEV_REDSHIFT_PORT'))
    },
}

# 飞书 Webhook URL
FEISHU_WEBHOOK_URL = 'web钩子'

# S3 配置
S3_BUCKET_NAME = 's3_namr'

# AWS 配置（请使用环境变量或配置文件替换）
AWS_ACCESS_KEY_ID = os.getenv('EVA_AWS_ACCESS_KEY_ID')
AWS_SECRET_ACCESS_KEY = os.getenv('EVA_AWS_SECRET_ACCESS_KEY')
AWS_REGION = os.getenv('EVA_REGION_NAME')

# 并发执行数量
MAX_CONCURRENT_TASKS = 15

# 是否在每个文件中异步执行 SQL 语句
ASYNC_PER_STATEMENT = True  # True 表示文件和 SQL 都异步，False 表示仅文件异步

# 已执行的文件和表名集合，防止重复执行
executed_files_and_tables: Set[str] = set()
executed_lock = asyncio.Lock()

# 失败的 SQL 语句列表
failed_statements: List[str] = []
failed_lock = asyncio.Lock()

# 异步数据库连接池
connection_pools: Dict[str, asyncpg.pool.Pool] = {}

# 在代码中指定 SQL 文件或目录的路径
SQL_PATH = '/opt/data/eva_data_warehouse/dwd_script/sql_day_script/insert_table_5.sql'


def load_sql_statements(path: str) -> List[Dict[str, Any]]:
    """
    从指定路径加载 SQL 语句，可以是目录或单个文件。
    返回包含文件名、目标表名和 SQL 语句的列表。
    """
    sql_statements = []

    if os.path.isfile(path):
        # 如果输入的是文件，处理该文件
        sql_files = [path]
    elif os.path.isdir(path):
        # 如果输入的是目录，处理目录中的所有 .sql 文件
        sql_files = glob.glob(os.path.join(path, '*.sql'))
    else:
        logging.error(f"指定的路径 '{path}' 既不是文件也不是目录。")
        return sql_statements  # 返回空列表

    for filepath in sql_files:
        with open(filepath, 'r', encoding='utf-8') as file:
            sql_content = file.read().strip()
            if sql_content:
                # 提取表名，可以根据实际情况修改解析方式
                table_name = extract_table_name(sql_content)
                if not table_name:
                    logging.warning(f"文件 {filepath} 中未找到目标表名，跳过该文件。")
                    continue
                sql_statements.append({
                    'file_path': filepath,
                    'table_name': table_name,
                    'sql': sql_content
                })
    if sql_statements:
        logging.info(f"从路径 {path} 加载了 {len(sql_statements)} 份 SQL 文件语句。")
    else:
        logging.info(f"路径 {path} 中未找到需要执行的 SQL 文件。")
    return sql_statements


def extract_table_name(sql_content: str) -> str:
    """
    使用 sqlparse 库从 SQL 语句中提取目标表名。
    """
    parsed = sqlparse.parse(sql_content)
    for statement in parsed:
        tokens = [token for token in statement.flatten() if not token.is_whitespace]
        for i, token in enumerate(tokens):
            if token.ttype == sqlparse.tokens.DML:
                keyword = token.value.upper()
                if keyword == 'SELECT':
                    for j in range(i + 1, len(tokens)):
                        if tokens[j].ttype == sqlparse.tokens.Keyword and tokens[j].value.upper() == 'FROM':
                            if j + 1 < len(tokens):
                                table_token = tokens[j + 1]
                                if table_token.ttype in (sqlparse.tokens.Name, sqlparse.tokens.Keyword):
                                    table_name = table_token.value.strip('"').split('.')[-1]
                                    return table_name
                            break
                elif keyword == 'INSERT':
                    if i + 2 < len(tokens):
                        next_token = tokens[i + 1]
                        table_token = tokens[i + 2]
                        if next_token.ttype == sqlparse.tokens.Keyword and next_token.value.upper() == 'INTO':
                            if table_token.ttype in (sqlparse.tokens.Name, sqlparse.tokens.Keyword):
                                table_name = table_token.value.strip('"').split('.')[-1]
                                return table_name
                elif keyword == 'UPDATE':
                    if i + 1 < len(tokens):
                        table_token = tokens[i + 1]
                        if table_token.ttype in (sqlparse.tokens.Name, sqlparse.tokens.Keyword):
                            table_name = table_token.value.strip('"').split('.')[-1]
                            return table_name
                elif keyword == 'DELETE':
                    if i + 2 < len(tokens):
                        next_token = tokens[i + 1]
                        table_token = tokens[i + 2]
                        if next_token.ttype == sqlparse.tokens.Keyword and next_token.value.upper() == 'FROM':
                            if table_token.ttype in (sqlparse.tokens.Name, sqlparse.tokens.Keyword):
                                table_name = table_token.value.strip('"').split('.')[-1]
                                return table_name
    return ""


@asynccontextmanager
async def get_connection(db_config_name: str):
    pool = connection_pools[db_config_name]
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)


async def execute_sql(file_info: Dict[str, Any], db_config_name: str) -> bool:
    """
    异步执行单个 SQL 文件中的所有语句。
    包含错误处理，防止重复执行。
    使用环境名称、文件名和表名的组合作为唯一键，防止重复执行相同的脚本。
    """
    file_path = file_info['file_path']
    table_name = file_info['table_name']
    sql_content = file_info['sql']

    unique_key = f"{db_config_name}:{file_path}:{table_name}"
    async with executed_lock:
        if unique_key in executed_files_and_tables:
            logging.info(f"[{db_config_name}] 文件 '{file_path}' 中的表 '{table_name}' 已在当前环境执行，跳过执行。")
            return True
        else:
            executed_files_and_tables.add(unique_key)

    # 使用 sqlparse 将 SQL 内容拆分为独立的语句
    statements = sqlparse.split(sql_content)
    success = True  # 标记该文件中是否所有语句都成功执行

    try:
        if ASYNC_PER_STATEMENT:
            # 模式1：文件和 SQL 都异步执行
            # 新增：创建信号量，限制文件内部的并发语句数量，防止过多的并发任务
            semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)

            async def execute_statement(stmt):
                nonlocal success
                stmt = stmt.strip()
                if stmt:
                    async with semaphore:
                        # 每个语句获取一个新的连接，确保线程安全
                        async with get_connection(db_config_name) as conn:
                            try:
                                await conn.execute(stmt)
                                logging.info(f"[{db_config_name}] 文件 '{file_path}' 中的语句执行成功：{stmt[:50]}...")
                            except Exception as e:
                                success = False
                                async with failed_lock:
                                    failed_statements.append(f"-- File: {file_path}\n{stmt}")
                                logging.exception(f"[{db_config_name}] 执行文件 '{file_path}' 中的语句时失败：{stmt[:50]}...")

            # 为每个语句创建异步任务
            tasks = [execute_statement(stmt) for stmt in statements]
            await asyncio.gather(*tasks)
        else:
            # 模式2：仅文件异步，SQL 语句按顺序执行
            async with get_connection(db_config_name) as conn:
                for stmt in statements:
                    stmt = stmt.strip()
                    if stmt:  # 确保语句不为空
                        try:
                            await conn.execute(stmt)
                            logging.info(f"[{db_config_name}] 文件 '{file_path}' 中的语句执行成功：{stmt[:50]}...")
                        except Exception as e:
                            success = False
                            async with failed_lock:
                                failed_statements.append(f"-- File: {file_path}\n{stmt}")
                            logging.exception(f"[{db_config_name}] 执行文件 '{file_path}' 中的语句时失败：{stmt[:50]}...")
                            # 如果希望继续执行后续语句，请注释掉下面的 break
                            # break

        return success
    except Exception as e:
        async with executed_lock:
            executed_files_and_tables.remove(unique_key)
        logging.exception(f"[{db_config_name}] 执行文件 '{file_path}' 时发生异常。")
        return False


async def upload_failed_sql_to_s3(failed_statements: List[str]) -> str:
    """
    将失败的 SQL 语句上传到 S3，并返回文件的访问链接。
    """
    # 生成时间戳，格式为 YYYYMMDD_HHMMSS
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    # 定义 S3 路径和文件名
    s3_key = f'failed_sql_statements_{timestamp}.sql'
    s3 = boto3.client(
        's3',
        region_name=AWS_REGION,
        aws_access_key_id=AWS_ACCESS_KEY_ID,
        aws_secret_access_key=AWS_SECRET_ACCESS_KEY
    )

    sql_content = ';\n'.join(failed_statements) + ';'

    try:
        s3.put_object(Bucket=S3_BUCKET_NAME, Key=s3_key, Body=sql_content.encode('utf-8'))
        logging.info(f"失败的 SQL 语句已上传到 S3：s3://{S3_BUCKET_NAME}/{s3_key}")
        # 构建文件的访问链接
        file_url = f"https://{S3_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{s3_key}"
        return file_url
    except ClientError as e:
        logging.exception(f"上传失败的 SQL 到 S3 时发生错误：{e}")
        return ""


def send_feishu_notification(file_url: str, db_config_name: str, failed_statements: List[str]):
    """
    通过飞书机器人发送通知，包含失败的 SQL 文件链接或直接附上失败的 SQL 语句。
    """
    if file_url:
        message = f"在执行 SQL 脚本到 {db_config_name} 环境时发生错误。\n" \
                  f"失败的 SQL 语句已上传到 S3：\n{file_url}\n请点击链接下载并查看。"
    else:
        if len(failed_statements) <= 5:
            failed_sql_content = '\n\n'.join(failed_statements)
            message = f"在执行 SQL 脚本到 {db_config_name} 环境时发生错误，且无法上传到 S3。\n" \
                      f"以下是失败的 SQL 语句：\n{failed_sql_content}\n请检查上传权限并处理。"
        else:
            message = f"在执行 SQL 脚本到 {db_config_name} 环境时发生错误，且无法上传到 S3。\n" \
                      f"失败的 SQL 语句较多，请检查上传权限并处理。"

    payload = {
        "msg_type": "text",
        "content": {
            "text": message
        }
    }
    try:
        response = requests.post(FEISHU_WEBHOOK_URL, json=payload, timeout=10)
        if response.status_code == 200:
            logging.info("飞书通知发送成功。")
        else:
            logging.error(f"飞书通知发送失败，状态码：{response.status_code}，响应内容：{response.text}")
    except Exception as e:
        logging.exception(f"发送飞书通知时发生异常：{e}")


async def main():
    # 在 main 方法中指定数据库环境
    db_config_name = 'dev'

    # 初始化数据库连接池
    if db_config_name not in DB_CONFIGS:
        logging.error(f"未知的数据库配置：{db_config_name}")
        return

    db_config = DB_CONFIGS[db_config_name]

    # 创建异步连接池，调整 max_size 以支持更多的并发连接
    pool = await asyncpg.create_pool(
        **db_config,
        min_size=1,
        max_size=MAX_CONCURRENT_TASKS * 2  # 设置为并发任务数量的倍数
    )
    connection_pools[db_config_name] = pool

    try:
        # 从指定路径加载 SQL 语句
        sql_statements_info = load_sql_statements(SQL_PATH)

        if not sql_statements_info:
            logging.info("未找到需要执行的 SQL 文件，程序结束。")
            return

        # 创建信号量，限制全局并发任务数量
        semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)

        async def sem_execute(file_info):
            async with semaphore:
                await execute_sql(file_info, db_config_name)

        # 异步并发执行 SQL 文件
        tasks = [sem_execute(info) for info in sql_statements_info]
        await asyncio.gather(*tasks)

        # 所有任务完成后，检查是否有失败的 SQL 语句
        if failed_statements:
            # 上传失败的 SQL 语句到 S3
            file_url = await upload_failed_sql_to_s3(failed_statements)
            # 发送飞书通知
            send_feishu_notification(file_url, db_config_name, failed_statements)
        else:
            logging.info("所有 SQL 语句执行成功。")
    finally:
        # 关闭数据库连接池
        await pool.close()


if __name__ == "__main__":
    asyncio.run(main())
